{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "30c61e75",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd \n",
    "import numpy as np\n",
    "import pysam \n",
    "import argparse\n",
    "from pybedtools import BedTool\n",
    "from pathlib import Path\n",
    "from io import StringIO\n",
    "from functools import reduce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a11cf833",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load arguments \n",
    "ge_matrix_flat_path = \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/misexpression_v2/2_misexp_qc/misexp_gene_cov_corr/tpm_zscore_4568smpls_8610genes_flat_misexp_corr_qc.csv\"\n",
    "chrom = \"chr21\"\n",
    "sv_info_path = \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/lof_missense/data/sv_vcf/info_table/final_sites_critical_info_allele.txt\"\n",
    "vcf_path=\"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/lof_missense/data/sv_vcf/filtered_merged_gs_svp_10728.vcf.gz\"\n",
    "gencode_path= \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/misexpression_v2/reference/gencode/gencode.v31.annotation.sorted.gtf.gz\"\n",
    "wgs_rna_paired_smpls_path = \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/misexpression_v2/1_rna_seq_qc/wgs_rna_match/paired_wgs_rna_postqc_prioritise_wgs.tsv\"\n",
    "z_cutoff_list = [2, 10, 20, 30, 40, 50, 60]\n",
    "tpm_cutoff = 0.5\n",
    "window_start = 0\n",
    "window_step_size = 200000\n",
    "window_max = 0\n",
    "af_lower, af_upper = (0, 0.01)\n",
    "vep_msc_path = \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/misexpression_v2/4_vrnt_enrich/sv_vep/msc/SV_vep_hg38_msc_parsed.tsv\"\n",
    "root_dir = \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/misexpression_v2/4_vrnt_enrich/sv_count_carriers/windows/200kb_window\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8348c7c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inputs:\n",
      "- Chromosome: chr21\n",
      "- AF range: 0-0.01\n",
      "- Window start: 0\n",
      "- Window size: 200000\n",
      "- Max window: 0\n",
      "- Z-score cutoffs: [2, 10, 20, 30, 40, 50, 60]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"Inputs:\")\n",
    "print(f\"- Chromosome: {chrom}\")\n",
    "print(f\"- AF range: {af_lower}-{af_upper}\")\n",
    "print(f\"- Window start: {window_start}\")\n",
    "print(f\"- Window size: {window_step_size}\")\n",
    "print(f\"- Max window: {window_max}\")\n",
    "print(f\"- Z-score cutoffs: {z_cutoff_list}\")\n",
    "print(\"\")\n",
    "\n",
    "# create root directory \n",
    "root_dir_path = Path(root_dir)\n",
    "root_dir_path.mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ab34fc3e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gene IDs passing filters: 8610\n",
      "RNA-seq sample IDs passing QC: 4568\n",
      "\n",
      "Loading input VCF ...\n",
      "VCF loaded.\n",
      "Subset VCF to samples with RNA-seq ...\n",
      "Number of samples in VCF with RNA ID and passing QC: 2640\n",
      "Number of RNA IDs passing QC: 2640\n"
     ]
    }
   ],
   "source": [
    "### Collect samples with VCF calls and RNA sequencing \n",
    "# read in gene expression file \n",
    "ge_matrix_flat_df = pd.read_csv(ge_matrix_flat_path, sep=\",\")\n",
    "gene_id_pass_qc_set = set(ge_matrix_flat_df.gene_id.unique())\n",
    "print(f\"Gene IDs passing filters: {len(gene_id_pass_qc_set)}\")\n",
    "smpl_id_pass_qc_set = set(ge_matrix_flat_df.rna_id.unique())\n",
    "print(f\"RNA-seq sample IDs passing QC: {len(smpl_id_pass_qc_set)}\")\n",
    "print(\"\")\n",
    "# egan ID, RNA ID sample links \n",
    "wgs_rna_paired_smpls_df = pd.read_csv(wgs_rna_paired_smpls_path, sep=\"\\t\")\n",
    "egan_ids_with_rna = wgs_rna_paired_smpls_df[wgs_rna_paired_smpls_df.rna_id.isin(smpl_id_pass_qc_set)].egan_id.tolist()\n",
    "# load VCF and subset to EGAN IDs with RNA \n",
    "print(\"Loading input VCF ...\")\n",
    "vcf_path = vcf_path\n",
    "vcf = pysam.VariantFile(vcf_path, mode = \"r\")\n",
    "print(\"VCF loaded.\")\n",
    "print(\"Subset VCF to samples with RNA-seq ...\")\n",
    "vcf_samples = [sample for sample in vcf.header.samples]\n",
    "vcf_egan_ids_with_rna = set(egan_ids_with_rna).intersection(set(vcf_samples))\n",
    "vcf.subset_samples(vcf_egan_ids_with_rna)\n",
    "vcf_samples_with_rna = [sample for sample in vcf.header.samples]\n",
    "print(f\"Number of samples in VCF with RNA ID and passing QC: {len(vcf_samples_with_rna)}\")\n",
    "# subset egan ID and RNA ID links to samples with SV calls and passing QC \n",
    "wgs_rna_paired_smpls_with_sv_calls_df = wgs_rna_paired_smpls_df[wgs_rna_paired_smpls_df.egan_id.isin(vcf_samples_with_rna)]\n",
    "# write EGAN-RNA ID pairs to file\n",
    "egan_rna_smpls_dir = root_dir_path.joinpath(\"egan_rna_smpls\")\n",
    "Path(egan_rna_smpls_dir).mkdir(parents=True, exist_ok=True)\n",
    "wgs_rna_paired_smpls_with_sv_calls_df.to_csv(egan_rna_smpls_dir.joinpath(\"egan_rna_ids_paired_pass_qc.tsv\"), sep=\"\\t\", index=False)\n",
    "rna_id_pass_qc_sv_calls = wgs_rna_paired_smpls_with_sv_calls_df.rna_id.unique().tolist()\n",
    "print(f\"Number of RNA IDs passing QC: {len(rna_id_pass_qc_sv_calls)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6465df21",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of genes pass QC on chr21: 148\n"
     ]
    }
   ],
   "source": [
    "### write bed file for genes on chromosome passing QC \n",
    "gene_bed_dir = root_dir_path.joinpath(\"genes_bed\")\n",
    "gene_bed_dir.mkdir(parents=True, exist_ok=True)\n",
    "gene_bed_path = gene_bed_dir.joinpath(f\"{chrom}_genes.bed\")\n",
    "gene_id_pass_qc_on_chrom = []\n",
    "with open(gene_bed_path, \"w\") as f:\n",
    "    for gtf in pysam.TabixFile(gencode_path).fetch(chrom, parser = pysam.asGTF()):\n",
    "        if gtf.feature == \"gene\" and gtf.gene_id.split('.')[0] in gene_id_pass_qc_set:\n",
    "            # check for multiple entries with same name \n",
    "            gene_id_list, chrom_list, start_list, end_list, strand_list = [], [], [], [], []\n",
    "            gene_id_list.append(gtf.gene_id.split('.')[0])\n",
    "            chrom_list.append(gtf.contig)\n",
    "            start_list.append(gtf.start)\n",
    "            end_list.append((gtf.end))\n",
    "            strand_list.append((gtf.strand))\n",
    "            # check or write output\n",
    "            if len(gene_id_list) > 1 or len(chrom_list) > 1 or len(start_list) > 1 or len(end_list) > 1:\n",
    "                print(f\"{gene_id} has multiple entries in gencode file - excluded from output file.\")\n",
    "            elif len(chrom_list) == 0 or len(start_list) == 0 or len(end_list) == 0: \n",
    "                print(f\"{gene_id} has no entries in gencode file - excluded from output file.\")\n",
    "            else: \n",
    "                gene_id, gtf_chrom, start, end, strand = gene_id_list[0], chrom_list[0], start_list[0], end_list[0], strand_list[0]\n",
    "                gene_id_pass_qc_on_chrom.append(gene_id)\n",
    "                chrom_num = gtf_chrom.split(\"chr\")[1]\n",
    "                f.write(f\"{chrom_num}\\t{start - 1}\\t{end}\\t{gene_id}\\t0\\t{strand}\\n\")\n",
    "print(f\"Number of genes pass QC on {chrom}: {len(gene_id_pass_qc_on_chrom)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9309e46e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sample IDs in gene expression matrix with SV calls: 2640\n"
     ]
    }
   ],
   "source": [
    "# subset gene expression file to genes on chromosome \n",
    "ge_matrix_flat_chrom_df = ge_matrix_flat_df[ge_matrix_flat_df.gene_id.isin(gene_id_pass_qc_on_chrom)]\n",
    "# subset gene expression file to samples with SV calls \n",
    "ge_matrix_flat_chrom_egan_df = pd.merge(ge_matrix_flat_chrom_df, wgs_rna_paired_smpls_with_sv_calls_df, how=\"inner\", on=\"rna_id\")\n",
    "print(f\"Sample IDs in gene expression matrix with SV calls: {len(ge_matrix_flat_chrom_egan_df.egan_id.unique())}\")\n",
    "# write to file \n",
    "ge_matrix_flat_chrom_egan_dir = root_dir_path.joinpath(\"express_mtx\")\n",
    "Path(ge_matrix_flat_chrom_egan_dir).mkdir(parents=True, exist_ok=True)\n",
    "ge_matrix_flat_chrom_egan_df.to_csv(ge_matrix_flat_chrom_egan_dir.joinpath(f\"{chrom}_ge_mtx_flat.tsv\"), sep=\"\\t\", index=False)\n",
    "# add gene-sample pair to expression dataframe\n",
    "ge_matrix_flat_chrom_egan_df[\"gene_smpl_pair\"] = ge_matrix_flat_chrom_egan_df.gene_id + \",\" + ge_matrix_flat_chrom_egan_df.rna_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "382d8f0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "### write SVs on chromosome to bed file \n",
    "vrnts_bed_dir = root_dir_path.joinpath(\"vrnts_bed\")\n",
    "vrnts_bed_dir.mkdir(parents=True, exist_ok=True)\n",
    "vrnts_bed_path = vrnts_bed_dir.joinpath(f\"{chrom}_vrnts.bed\")\n",
    "with open(sv_info_path, \"r\") as f_in, open(vrnts_bed_path, \"w\") as f_out:\n",
    "    for line in f_in:\n",
    "        if line.startswith(\"plinkID\"): \n",
    "            continue\n",
    "        else: \n",
    "            vrnt_id, sv_chrom, pos, end = line.split(\"\\t\")[0], line.split(\"\\t\")[2], line.split(\"\\t\")[3], line.split(\"\\t\")[4]\n",
    "            if sv_chrom == chrom:\n",
    "                chrom_num = sv_chrom.split(\"chr\")[1]\n",
    "                # convert to zero-based\n",
    "                start = int(pos) - 1 \n",
    "                f_out.write(f\"{chrom_num}\\t{start}\\t{end}\\t{vrnt_id}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5c2e35c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SV types in SV info file: ['DEL', 'MEI', 'DUP', 'INV']\n"
     ]
    }
   ],
   "source": [
    "### SV information \n",
    "sv_info_df = pd.read_csv(sv_info_path, sep=\"\\t\", dtype={\"plinkID\":str})\n",
    "sv_types_list = sv_info_df.SVTYPE.unique().tolist()\n",
    "print(f\"SV types in SV info file: {sv_types_list}\")\n",
    "sv_info_id_af_df = sv_info_df[[\"plinkID\", \"AF\", \"SVTYPE\"]].rename(columns={\"plinkID\":\"vrnt_id\"})\n",
    "\n",
    "### generate output directories \n",
    "# variant-window intersection directory \n",
    "intersect_bed_dir=root_dir_path.joinpath(\"intersect_bed\")\n",
    "Path(intersect_bed_dir).mkdir(parents=True, exist_ok=True)\n",
    "# genotypes of variants in windows directory \n",
    "intersect_vrnt_gts_dir = root_dir_path.joinpath(\"intersect_vrnt_gts\")\n",
    "Path(intersect_vrnt_gts_dir).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# load variant and gene bed files \n",
    "vrnts_bed = BedTool(vrnts_bed_path)\n",
    "genes_bed = BedTool(gene_bed_path)\n",
    "# column names for variant window intersection bed file \n",
    "intersect_bed_columns={0:\"chrom_gene\", 1:\"start_gene\", 2:\"end_gene\", 3: \"gene_id\", 4: \"score\", 5: \"strand\", \n",
    "                       6:\"chrom_vrnt\", 7:\"start_vrnt\", 8:\"end_vrnt\", 9:\"vrnt_id\"}\n",
    "\n",
    "# load most severe consequence \n",
    "vep_msc_df = pd.read_csv(vep_msc_path, sep=\"\\t\", dtype={\"vrnt_id\": str})\n",
    "msc_list = vep_msc_df.Consequence.unique().tolist()\n",
    "vep_msc_cnsqn_df = vep_msc_df.rename(columns={\"Uploaded_variation\": \"vrnt_id\"})[[\"vrnt_id\", \"Consequence\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "64be6066",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counting variants upstream, 0bp ...\n",
      "Counting variants downstream, 0bp ...\n"
     ]
    }
   ],
   "source": [
    "entry = 0 \n",
    "carrier_count = {}\n",
    "for direction in [\"upstream\", \"downstream\"]:\n",
    "    window_size = window_start\n",
    "    seen_sv_gene_pairs = set()\n",
    "    while window_size <= window_max:\n",
    "        print(f\"Counting variants {direction}, {window_size}bp ...\")\n",
    "        if direction == \"upstream\": \n",
    "            l, r = window_size, 0 \n",
    "        else: \n",
    "            l, r = 0, window_size\n",
    "        vrnts_window_intersect_str = StringIO(str(genes_bed.window(vrnts_bed, \n",
    "                                                                   l=l, \n",
    "                                                                   r=r,\n",
    "                                                                   sw=True\n",
    "                                                                  )))        \n",
    "        intersect_bed_df = pd.read_csv(vrnts_window_intersect_str, sep=\"\\t\", header=None).rename(columns=intersect_bed_columns)\n",
    "        # variant gene ID column \n",
    "        intersect_bed_df[\"vrnt_gene_id\"] = intersect_bed_df[\"vrnt_id\"].astype(str) + \",\" + intersect_bed_df[\"gene_id\"].astype(str)\n",
    "        # subset to variant gene pairs that have not been seen in previous windows    \n",
    "        sv_gene_pairs_in_window = set(intersect_bed_df.vrnt_gene_id.unique())\n",
    "        new_sv_gene_pairs = sv_gene_pairs_in_window - seen_sv_gene_pairs\n",
    "        if new_sv_gene_pairs.union(seen_sv_gene_pairs) != sv_gene_pairs_in_window: \n",
    "            raise ValueError(\"Missing variants from windows.\")\n",
    "        intersect_bed_new_sv_gene_df = intersect_bed_df[intersect_bed_df[\"vrnt_gene_id\"].isin(new_sv_gene_pairs)] \n",
    "        # generate set of variant IDs inside gene windows\n",
    "        intersect_vrnt_ids_no_dupl_df = intersect_bed_new_sv_gene_df[[\"chrom_vrnt\", \"vrnt_id\", \"start_vrnt\", \"end_vrnt\"]].drop_duplicates()\n",
    "        intersect_vrnt_ids_list = intersect_vrnt_ids_no_dupl_df.vrnt_id.unique()\n",
    "        # genotypes for variant IDs in windows\n",
    "        vrnt_gt_egan_dict = {}\n",
    "        count = 0\n",
    "        for vrnt_id in intersect_vrnt_ids_list:\n",
    "            # get chromosome, start and end of SV \n",
    "            chrom_vrnt, pos, end, = [intersect_vrnt_ids_no_dupl_df[intersect_vrnt_ids_no_dupl_df.vrnt_id == vrnt_id][col].item() for col in [\"chrom_vrnt\", \"start_vrnt\", \"end_vrnt\"]]\n",
    "            records = vcf.fetch(str(chrom_vrnt), pos, end)\n",
    "            found_vrnt_id = False\n",
    "            for rec in records: \n",
    "                vcf_vrnt_id = str(rec.id)\n",
    "                if vrnt_id == vcf_vrnt_id:\n",
    "                    found_vrnt_id = True \n",
    "                    gts = [s[\"GT\"] for s in rec.samples.values()]\n",
    "                    for i, gt in enumerate(gts): \n",
    "                        vrnt_gt_egan_dict[count] = [vrnt_id, vcf.header.samples[i], gt]\n",
    "                        count += 1 \n",
    "            if not found_vrnt_id: \n",
    "                raise ValueError(f\"Did not find {vrnt_id} in {vcf_path}\")\n",
    "        vrnt_gt_egan_nogene_df = pd.DataFrame.from_dict(vrnt_gt_egan_dict, orient=\"index\", columns=[\"vrnt_id\", \"egan_id\", \"genotype\"])\n",
    "        vrnt_gt_egan_nogene_df = vrnt_gt_egan_nogene_df.astype({\"genotype\":str})\n",
    "\n",
    "        # merge \n",
    "        dfs_to_merge = [vrnt_gt_egan_nogene_df, \n",
    "                        intersect_bed_new_sv_gene_df[[\"vrnt_id\", \"gene_id\",\"vrnt_gene_id\"]].drop_duplicates(),\n",
    "                        sv_info_id_af_df, \n",
    "                        vep_msc_cnsqn_df\n",
    "                       ]\n",
    "        df_merged = reduce(lambda  left,right: pd.merge(left,right,on=['vrnt_id'],\n",
    "                                                    how='inner'), dfs_to_merge)\n",
    "        # add expression \n",
    "        sv_intersect_express_info_df = pd.merge(df_merged, \n",
    "                                             ge_matrix_flat_chrom_egan_df,                  \n",
    "                                             how=\"inner\",\n",
    "                                             on=[\"egan_id\", \"gene_id\"])\n",
    "        # write to file \n",
    "        intersect_vrnt_gts_path = intersect_vrnt_gts_dir.joinpath(f\"{chrom}_intersect_vrnt_gts_{direction}_{window_size}.tsv\")\n",
    "        sv_intersect_express_info_df.to_csv(intersect_vrnt_gts_path, sep=\"\\t\", index=False)\n",
    "        \n",
    "        ### count carriers in test and controls across different Z-score cutoffs \n",
    "        for z_cutoff in z_cutoff_list: \n",
    "            carrier_count[entry] = [chrom, direction, window_size, z_cutoff]\n",
    "            ### get control and test gene-sample pairs based on TPM cutoff \n",
    "\n",
    "            # get gene-sample pairs passing z-score and TPM cutoff \n",
    "            misexp_df = ge_matrix_flat_chrom_egan_df[(ge_matrix_flat_chrom_egan_df[\"z-score\"] > z_cutoff) &\n",
    "                                                     (ge_matrix_flat_chrom_egan_df[\"TPM\"] > tpm_cutoff)\n",
    "                                                    ]\n",
    "            test_gene_ids = misexp_df.gene_id.unique()\n",
    "            test_gene_smpl_pairs = misexp_df.gene_smpl_pair.unique()\n",
    "            # get gene-sample pairs in control group - same gene set, samples not misexpressing\n",
    "            cntrl_gene_smpl_pairs = ge_matrix_flat_chrom_egan_df[~ge_matrix_flat_chrom_egan_df.gene_smpl_pair.isin(test_gene_smpl_pairs) &\n",
    "                                                                 ge_matrix_flat_chrom_egan_df.gene_id.isin(test_gene_ids)\n",
    "                                                                ].gene_smpl_pair.unique()\n",
    "            # check overlap between sets is empty \n",
    "            if len(set(test_gene_smpl_pairs).intersection(set(cntrl_gene_smpl_pairs))) != 0:\n",
    "                raise ValueError(\"Overlap between control and test gene-sample pair sets\")\n",
    "            # check number of test and control gene-pairs matches expected \n",
    "            total_test_control_pairs = len(test_gene_smpl_pairs) + len(cntrl_gene_smpl_pairs)\n",
    "            test_genes_by_smpls = len(test_gene_ids) * len(rna_id_pass_qc_sv_calls)\n",
    "            if total_test_control_pairs !=  test_genes_by_smpls: \n",
    "                raise ValueError(\"Number of test and control gene pairs does not match test gene IDs by samples\")\n",
    "            carrier_count[entry] += [len(test_gene_ids), len(rna_id_pass_qc_sv_calls), len(test_gene_smpl_pairs), len(cntrl_gene_smpl_pairs)]\n",
    "\n",
    "            ### count carriers in control and test groups \n",
    "            # create column with gene-sample pairs  \n",
    "            sv_intersect_express_info_df[\"gene_smpls_id\"] = sv_intersect_express_info_df.gene_id.astype(str) + \",\" + sv_intersect_express_info_df.rna_id.astype(str)\n",
    "            # get misexpression carriers \n",
    "            misexpress_carriers_df = sv_intersect_express_info_df[(sv_intersect_express_info_df.gene_smpls_id.isin(test_gene_smpl_pairs)) & \n",
    "                                                                  (sv_intersect_express_info_df.AF >= af_lower) & \n",
    "                                                                  (sv_intersect_express_info_df.AF < af_upper) &\n",
    "                                                                  (sv_intersect_express_info_df.genotype.isin(['(0, 1)', '(1, 1)']))\n",
    "                                                                 ].copy()\n",
    "            # get non-misexpression carriers\n",
    "            cntrl_carriers_df = sv_intersect_express_info_df[(sv_intersect_express_info_df.gene_smpls_id.isin(cntrl_gene_smpl_pairs)) &\n",
    "                                                                      (sv_intersect_express_info_df.AF >= af_lower) & \n",
    "                                                                      (sv_intersect_express_info_df.AF < af_upper) &\n",
    "                                                                      (sv_intersect_express_info_df.genotype.isin(['(0, 1)', '(1, 1)']))\n",
    "                                                                     ].copy()\n",
    "            # add number of carriers \n",
    "            gene_smpl_misexpress_carriers = misexpress_carriers_df.gene_smpls_id.unique()\n",
    "            gene_smpl_cntrl_carriers = cntrl_carriers_df.gene_smpls_id.unique()\n",
    "            carrier_count[entry] += [len(gene_smpl_misexpress_carriers), len(gene_smpl_cntrl_carriers)]\n",
    "            # check overlap between sets is empty \n",
    "            if len(set(gene_smpl_misexpress_carriers).intersection(set(gene_smpl_cntrl_carriers))) != 0:\n",
    "                raise ValueError(\"Overlap between control and test gene-sample pair carrier sets\")\n",
    "\n",
    "            ### count carriers for SV types and most severe consequence \n",
    "            # question remains whether to assign priority to different SV classes\n",
    "            for sv_type in sv_types_list:\n",
    "                misexpress_carriers_sv_type_df = misexpress_carriers_df[misexpress_carriers_df.SVTYPE == sv_type]\n",
    "                cntrl_carriers_sv_type_df = cntrl_carriers_df[cntrl_carriers_df.SVTYPE == sv_type]\n",
    "                sv_type_misexp_carriers = len(misexpress_carriers_sv_type_df.gene_smpls_id.unique())\n",
    "                sv_type_cntrl_carriers = len(cntrl_carriers_sv_type_df.gene_smpls_id.unique())\n",
    "                carrier_count[entry] += [sv_type_misexp_carriers, sv_type_cntrl_carriers]\n",
    "                for msc in msc_list: \n",
    "                    msc_type_misexp_carriers = len(misexpress_carriers_sv_type_df[misexpress_carriers_sv_type_df.Consequence == msc].gene_smpls_id.unique())\n",
    "                    msc_type_cntrl_carriers = len(cntrl_carriers_sv_type_df[cntrl_carriers_sv_type_df.Consequence == msc].gene_smpls_id.unique())\n",
    "                    carrier_count[entry] += [msc_type_misexp_carriers, msc_type_cntrl_carriers]\n",
    "            entry += 1\n",
    "        # add variant gene pairs to seen set \n",
    "        seen_sv_gene_pairs = seen_sv_gene_pairs.union(sv_gene_pairs_in_window)\n",
    "        window_size += window_step_size\n",
    "\n",
    "# write results \n",
    "vrnt_type_cols = []\n",
    "for sv_type in sv_types_list: \n",
    "    vrnt_type_cols += [f\"{sv_type}_misexp\", f\"{sv_type}_contrl\"]\n",
    "    for msc in msc_list:\n",
    "        vrnt_type_cols += [f\"{sv_type}_{msc}_misexp\", f\"{sv_type}_{msc}_contrl\"]\n",
    "\n",
    "columns = [\"chrom\", \"direction\", \"window_size\", \"z_cutoff\", \"misexp_genes\", \"smpls_pass_qc\", \"total_misexp\", \"total_control\", \"all_sv_misexp\", \"all_sv_contrl\"] + vrnt_type_cols\n",
    "carrier_count_df = pd.DataFrame.from_dict(carrier_count, orient=\"index\", columns=columns)\n",
    "#carrier_count_dir = f\"{root_dir}/carrier_count\"\n",
    "#Path(carrier_count_dir).mkdir(parents=True, exist_ok=True)\n",
    "#carrier_count_path = f\"{carrier_count_dir}/{chrom}_carrier_count.tsv\"\n",
    "#carrier_count_df.to_csv(carrier_count_path, sep=\"\\t\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "1bac0f99-57f7-460e-af14-e73b7e34421c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>chrom_gene</th>\n",
       "      <th>start_gene</th>\n",
       "      <th>end_gene</th>\n",
       "      <th>gene_id</th>\n",
       "      <th>score</th>\n",
       "      <th>strand</th>\n",
       "      <th>chrom_vrnt</th>\n",
       "      <th>start_vrnt</th>\n",
       "      <th>end_vrnt</th>\n",
       "      <th>vrnt_id</th>\n",
       "      <th>vrnt_gene_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>21</td>\n",
       "      <td>10521551</td>\n",
       "      <td>10606140</td>\n",
       "      <td>ENSG00000274391</td>\n",
       "      <td>0</td>\n",
       "      <td>+</td>\n",
       "      <td>21</td>\n",
       "      <td>10571600</td>\n",
       "      <td>10573190</td>\n",
       "      <td>241500</td>\n",
       "      <td>241500,ENSG00000274391</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>21</td>\n",
       "      <td>10521551</td>\n",
       "      <td>10606140</td>\n",
       "      <td>ENSG00000274391</td>\n",
       "      <td>0</td>\n",
       "      <td>+</td>\n",
       "      <td>21</td>\n",
       "      <td>10572648</td>\n",
       "      <td>10577358</td>\n",
       "      <td>DEL_chr21_10572649_10577358</td>\n",
       "      <td>DEL_chr21_10572649_10577358,ENSG00000274391</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>21</td>\n",
       "      <td>13516105</td>\n",
       "      <td>13574999</td>\n",
       "      <td>ENSG00000228159</td>\n",
       "      <td>0</td>\n",
       "      <td>+</td>\n",
       "      <td>21</td>\n",
       "      <td>13563163</td>\n",
       "      <td>13568338</td>\n",
       "      <td>241763</td>\n",
       "      <td>241763,ENSG00000228159</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>21</td>\n",
       "      <td>13516105</td>\n",
       "      <td>13574999</td>\n",
       "      <td>ENSG00000228159</td>\n",
       "      <td>0</td>\n",
       "      <td>+</td>\n",
       "      <td>21</td>\n",
       "      <td>13563152</td>\n",
       "      <td>13568328</td>\n",
       "      <td>DEL_chr21_13563153_13568328</td>\n",
       "      <td>DEL_chr21_13563153_13568328,ENSG00000228159</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>21</td>\n",
       "      <td>14108811</td>\n",
       "      <td>14210891</td>\n",
       "      <td>ENSG00000188992</td>\n",
       "      <td>0</td>\n",
       "      <td>-</td>\n",
       "      <td>21</td>\n",
       "      <td>14173713</td>\n",
       "      <td>14173907</td>\n",
       "      <td>241894</td>\n",
       "      <td>241894,ENSG00000188992</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>280</th>\n",
       "      <td>21</td>\n",
       "      <td>46037050</td>\n",
       "      <td>46039807</td>\n",
       "      <td>ENSG00000224413</td>\n",
       "      <td>0</td>\n",
       "      <td>+</td>\n",
       "      <td>21</td>\n",
       "      <td>46028122</td>\n",
       "      <td>46328545</td>\n",
       "      <td>245161</td>\n",
       "      <td>245161,ENSG00000224413</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>281</th>\n",
       "      <td>21</td>\n",
       "      <td>46037050</td>\n",
       "      <td>46039807</td>\n",
       "      <td>ENSG00000224413</td>\n",
       "      <td>0</td>\n",
       "      <td>+</td>\n",
       "      <td>21</td>\n",
       "      <td>46022145</td>\n",
       "      <td>46362649</td>\n",
       "      <td>414879</td>\n",
       "      <td>414879,ENSG00000224413</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>282</th>\n",
       "      <td>21</td>\n",
       "      <td>46151612</td>\n",
       "      <td>46152647</td>\n",
       "      <td>ENSG00000237338</td>\n",
       "      <td>0</td>\n",
       "      <td>+</td>\n",
       "      <td>21</td>\n",
       "      <td>46028122</td>\n",
       "      <td>46328545</td>\n",
       "      <td>245161</td>\n",
       "      <td>245161,ENSG00000237338</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>283</th>\n",
       "      <td>21</td>\n",
       "      <td>46151612</td>\n",
       "      <td>46152647</td>\n",
       "      <td>ENSG00000237338</td>\n",
       "      <td>0</td>\n",
       "      <td>+</td>\n",
       "      <td>21</td>\n",
       "      <td>46126039</td>\n",
       "      <td>46347637</td>\n",
       "      <td>245214</td>\n",
       "      <td>245214,ENSG00000237338</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>284</th>\n",
       "      <td>21</td>\n",
       "      <td>46151612</td>\n",
       "      <td>46152647</td>\n",
       "      <td>ENSG00000237338</td>\n",
       "      <td>0</td>\n",
       "      <td>+</td>\n",
       "      <td>21</td>\n",
       "      <td>46022145</td>\n",
       "      <td>46362649</td>\n",
       "      <td>414879</td>\n",
       "      <td>414879,ENSG00000237338</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>285 rows × 11 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     chrom_gene  start_gene  end_gene          gene_id  score strand  \\\n",
       "0            21    10521551  10606140  ENSG00000274391      0      +   \n",
       "1            21    10521551  10606140  ENSG00000274391      0      +   \n",
       "2            21    13516105  13574999  ENSG00000228159      0      +   \n",
       "3            21    13516105  13574999  ENSG00000228159      0      +   \n",
       "4            21    14108811  14210891  ENSG00000188992      0      -   \n",
       "..          ...         ...       ...              ...    ...    ...   \n",
       "280          21    46037050  46039807  ENSG00000224413      0      +   \n",
       "281          21    46037050  46039807  ENSG00000224413      0      +   \n",
       "282          21    46151612  46152647  ENSG00000237338      0      +   \n",
       "283          21    46151612  46152647  ENSG00000237338      0      +   \n",
       "284          21    46151612  46152647  ENSG00000237338      0      +   \n",
       "\n",
       "     chrom_vrnt  start_vrnt  end_vrnt                      vrnt_id  \\\n",
       "0            21    10571600  10573190                       241500   \n",
       "1            21    10572648  10577358  DEL_chr21_10572649_10577358   \n",
       "2            21    13563163  13568338                       241763   \n",
       "3            21    13563152  13568328  DEL_chr21_13563153_13568328   \n",
       "4            21    14173713  14173907                       241894   \n",
       "..          ...         ...       ...                          ...   \n",
       "280          21    46028122  46328545                       245161   \n",
       "281          21    46022145  46362649                       414879   \n",
       "282          21    46028122  46328545                       245161   \n",
       "283          21    46126039  46347637                       245214   \n",
       "284          21    46022145  46362649                       414879   \n",
       "\n",
       "                                    vrnt_gene_id  \n",
       "0                         241500,ENSG00000274391  \n",
       "1    DEL_chr21_10572649_10577358,ENSG00000274391  \n",
       "2                         241763,ENSG00000228159  \n",
       "3    DEL_chr21_13563153_13568328,ENSG00000228159  \n",
       "4                         241894,ENSG00000188992  \n",
       "..                                           ...  \n",
       "280                       245161,ENSG00000224413  \n",
       "281                       414879,ENSG00000224413  \n",
       "282                       245161,ENSG00000237338  \n",
       "283                       245214,ENSG00000237338  \n",
       "284                       414879,ENSG00000237338  \n",
       "\n",
       "[285 rows x 11 columns]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "intersect_bed_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e0a35f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "        ### count carriers in test and controls across different Z-score cutoffs \n",
    "        for z_cutoff in z_cutoff_list: \n",
    "            carrier_count[entry] = [chrom, direction, window_size, z_cutoff]\n",
    "            ### get control and test gene-sample pairs based on TPM cutoff \n",
    "\n",
    "            # get gene-sample pairs passing z-score and TPM cutoff \n",
    "            misexp_df = ge_matrix_flat_chrom_egan_df[(ge_matrix_flat_chrom_egan_df[\"z-score\"] > z_cutoff) &\n",
    "                                                          (ge_matrix_flat_chrom_egan_df[\"TPM\"] > tpm_cutoff)\n",
    "                                                         ]\n",
    "            test_gene_ids = misexp_df.gene_id.unique()\n",
    "            test_gene_smpl_pairs = misexp_df.gene_smpl_pair.unique()\n",
    "            # get gene-sample pairs in control group - same gene set, samples not misexpressing\n",
    "            cntrl_gene_smpl_pairs = ge_matrix_flat_chrom_egan_df[~ge_matrix_flat_chrom_egan_df.gene_smpl_pair.isin(test_gene_smpl_pairs) &\n",
    "                                                                 ge_matrix_flat_chrom_egan_df.gene_id.isin(test_gene_ids)\n",
    "                                                                ].gene_smpl_pair.unique()\n",
    "            # check overlap between sets is empty \n",
    "            if len(set(test_gene_smpl_pairs).intersection(set(cntrl_gene_smpl_pairs))) != 0:\n",
    "                raise ValueError(\"Overlap between control and test gene-sample pair sets\")\n",
    "            # check number of test and control gene-pairs matches expected \n",
    "            total_test_control_pairs = len(test_gene_smpl_pairs) + len(cntrl_gene_smpl_pairs)\n",
    "            test_genes_by_smpls = len(test_gene_ids) * len(rna_id_pass_qc_sv_calls)\n",
    "            if total_test_control_pairs !=  test_genes_by_smpls: \n",
    "                raise ValueError(\"Number of test and control gene pairs does not match test gene IDs by samples\")\n",
    "            carrier_count[entry] += [len(test_gene_ids), len(rna_id_pass_qc_sv_calls), len(test_gene_smpl_pairs), len(cntrl_gene_smpl_pairs)]\n",
    "\n",
    "            ### count carriers in control and test groups \n",
    "            # create column with gene-sample pairs  \n",
    "            sv_intersect_express_info_df[\"gene_smpls_id\"] = sv_intersect_express_info_df.gene_id.astype(str) + \",\" + sv_intersect_express_info_df.rna_id.astype(str)\n",
    "            # get misexpression carriers \n",
    "            misexpress_carriers_df = sv_intersect_express_info_df[(sv_intersect_express_info_df.gene_smpls_id.isin(test_gene_smpl_pairs)) & \n",
    "                                                                  (sv_intersect_express_info_df.AF >= af_lower) & \n",
    "                                                                  (sv_intersect_express_info_df.AF < af_upper) &\n",
    "                                                                  (sv_intersect_express_info_df.genotype.isin(['(0, 1)', '(1, 1)']))\n",
    "                                                                 ].copy()\n",
    "            # get non-misexpression carriers\n",
    "            cntrl_carriers_df = sv_intersect_express_info_df[(sv_intersect_express_info_df.gene_smpls_id.isin(cntrl_gene_smpl_pairs)) &\n",
    "                                                                      (sv_intersect_express_info_df.AF >= af_lower) & \n",
    "                                                                      (sv_intersect_express_info_df.AF < af_upper) &\n",
    "                                                                      (sv_intersect_express_info_df.genotype.isin(['(0, 1)', '(1, 1)']))\n",
    "                                                                     ].copy()\n",
    "            # add number of carriers \n",
    "            gene_smpl_misexpress_carriers = misexpress_carriers_df.gene_smpls_id.unique()\n",
    "            gene_smpl_cntrl_carriers = cntrl_carriers_df.gene_smpls_id.unique()\n",
    "            carrier_count[entry] += [len(gene_smpl_misexpress_carriers), len(gene_smpl_cntrl_carriers)]\n",
    "            # check overlap between sets is empty \n",
    "            if len(set(gene_smpl_misexpress_carriers).intersection(set(gene_smpl_cntrl_carriers))) != 0:\n",
    "                raise ValueError(\"Overlap between control and test gene-sample pair carrier sets\")\n",
    "\n",
    "            ### count carriers for SV types and most severe consequence \n",
    "            # question remains whether to assign priority to different SV classes\n",
    "            for sv_type in sv_types_list:\n",
    "                misexpress_carriers_sv_type_df = misexpress_carriers_df[misexpress_carriers_df.SVTYPE == sv_type]\n",
    "                cntrl_carriers_sv_type_df = cntrl_carriers_df[cntrl_carriers_df.SVTYPE == sv_type]\n",
    "                sv_type_misexp_carriers = len(misexpress_carriers_sv_type_df.gene_smpls_id.unique())\n",
    "                sv_type_cntrl_carriers = len(cntrl_carriers_sv_type_df.gene_smpls_id.unique())\n",
    "                carrier_count[entry] += [sv_type_misexp_carriers, sv_type_cntrl_carriers]\n",
    "                for msc in msc_list: \n",
    "                    msc_type_misexp_carriers = len(misexpress_carriers_sv_type_df[misexpress_carriers_sv_type_df.Consequence == msc].gene_smpls_id.unique())\n",
    "                    msc_type_cntrl_carriers = len(cntrl_carriers_sv_type_df[cntrl_carriers_sv_type_df.Consequence == msc].gene_smpls_id.unique())\n",
    "                    carrier_count[entry] += [msc_type_misexp_carriers, msc_type_cntrl_carriers]\n",
    "            entry += 1\n",
    "        # add variant gene pairs to seen set \n",
    "        seen_sv_gene_pairs = seen_sv_gene_pairs.union(sv_gene_pairs_in_window)\n",
    "        window_size += window_step_size\n",
    "\n",
    "# write results \n",
    "vrnt_type_cols = []\n",
    "for sv_type in sv_types_list: \n",
    "    vrnt_type_cols += [f\"{sv_type}_misexp\", f\"{sv_type}_contrl\"]\n",
    "    for msc in msc_list:\n",
    "        vrnt_type_cols += [f\"{sv_type}_{msc}_misexp\", f\"{sv_type}_{msc}_contrl\"]\n",
    "\n",
    "columns = [\"chrom\", \"direction\", \"window_size\", \"z_cutoff\", \"misexp_genes\", \"smpls_pass_qc\", \"total_misexp\", \"total_control\", \"all_sv_misexp\", \"all_sv_contrl\"] + vrnt_type_cols\n",
    "carrier_count_df = pd.DataFrame.from_dict(carrier_count, orient=\"index\", columns=columns)\n",
    "carrier_count_dir = f\"{root_dir}/carrier_count\"\n",
    "Path(carrier_count_dir).mkdir(parents=True, exist_ok=True)\n",
    "carrier_count_path = f\"{carrier_count_dir}/{chrom}_carrier_count.tsv\"\n",
    "carrier_count_df.to_csv(carrier_count_path, sep=\"\\t\", index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
